package org.apache.flink.runtime.accumulators;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import org.apache.flink.api.common.ConsumerID;
import org.apache.flink.api.common.accumulators.Accumulator;
import org.apache.flink.runtime.resourceconsumer.ConsumeAttemptID;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** Main accumulator registry which encapsulates user-defined accumulators. */
public class AccumulatorRegistry {

    protected static final Logger LOG = LoggerFactory.getLogger(AccumulatorRegistry.class);

    protected final ConsumerID consumerId;
    protected final ConsumeAttemptID taskID;

    /* User-defined Accumulator values stored for the executing task. */
    private final Map<String, Accumulator<?, ?>> userAccumulators = new ConcurrentHashMap<>(4);

    public AccumulatorRegistry(ConsumerID consumerId, ConsumeAttemptID taskID) {
        this.consumerId = consumerId;
        this.taskID = taskID;
    }

    /**
     * Creates a snapshot of this accumulator registry.
     *
     * @return a serialized accumulator map
     */
    public AccumulatorSnapshot getSnapshot() {
        try {
            return new AccumulatorSnapshot(consumerId, taskID, userAccumulators);
        } catch (Throwable e) {
            LOG.warn("Failed to serialize accumulators for task.", e);
            return null;
        }
    }

    /** Gets the map for user-defined accumulators. */
    public Map<String, Accumulator<?, ?>> getUserMap() {
        return userAccumulators;
    }
}
